#!/bin/sh

BR_IFACE="br-lan"
TC_URL="https://downloads.openwrt.org/releases/24.10.0/packages/aarch64_cortex-a53/base/tc-full_6.11.0-r1_aarch64_cortex-a53.ipk"

echo "==== KWRT Wi-Fi 限速管理脚本 ===="

# 检查 ip 命令
if ! command -v ip >/dev/null 2>&1; then
    echo "[错误] 未检测到 ip 命令，请先安装 ip-full 包"
    exit 1
fi

# 检查 tc 命令
if ! command -v tc >/dev/null 2>&1; then
    echo "[提示] 未检测到 tc 命令，正在自动安装 tc-full..."
    opkg update
    opkg install $TC_URL
    if [ $? -ne 0 ]; then
        echo "[错误] 安装 tc-full 失败，请检查网络或手动安装 IPK"
        exit 1
    fi
    echo "[完成] tc-full 已安装"
fi

# 获取当前 IPv4 设备列表
get_devices() {
    DEVICES=$(ip -4 neigh show dev $BR_IFACE | awk '{print $1" "$5}' | grep -v "FAILED")
}

# 显示设备并编号
show_devices() {
    echo "============================="
    echo "当前已连接 IPv4 设备列表："
    i=1
    echo "$DEVICES" | while read -r IP MAC; do
        echo "$i) IP:$IP  MAC:$MAC"
        i=$((i+1))
    done
    echo "============================="
}

# 显示限速规则
show_limits() {
    echo "当前生效的限速规则："
    tc -s class show dev $BR_IFACE | grep -v "root"
    echo "============================="
}

# 删除所有限速
delete_limits() {
    tc qdisc del dev $BR_IFACE root 2>/dev/null
    echo "[完成] 已删除所有限速规则"
}

# 主菜单循环
while true; do
    echo "-----------------------------"
    echo "KWRT Wi-Fi 限速管理"
    echo "1) 添加限速"
    echo "2) 删除所有限速"
    echo "3) 查看限速规则"
    echo "4) 退出"
    echo "-----------------------------"
    read -p "请选择操作 [1-4]: " choice

    case $choice in
        1)
            get_devices
            if [ -z "$DEVICES" ]; then
                echo "[提示] 未检测到任何设备，确认接口 $BR_IFACE 正确并有设备连接"
                continue
            fi
            show_devices
            read -p "输入要限制的设备序号 (多个序号用空格): " SELECTIONS
            if [ -z "$SELECTIONS" ]; then
                echo "[提示] 未输入序号，返回菜单"
                continue
            fi

            read -p "下载速率 (MB/s): " DL_MB
            read -p "上传速率 (MB/s): " UL_MB

            # 转换成 kbit
            DL_KBPS=$(echo "$DL_MB*1024*8" | bc | awk '{printf "%d\n",$1}')
            UL_KBPS=$(echo "$UL_MB*1024*8" | bc | awk '{printf "%d\n",$1}')

            # 创建根队列和父类（如果不存在）
            tc qdisc show dev $BR_IFACE | grep "htb" >/dev/null 2>&1
            if [ $? -ne 0 ]; then
                echo "[提示] 创建根队列和父类"
                tc qdisc add dev $BR_IFACE root handle 1: htb default 30
                tc class add dev $BR_IFACE parent 1: classid 1:1 htb rate 1gbit ceil 1gbit
            fi

            # 遍历用户选择的序号添加限速
            i=1
            echo "$DEVICES" | while read -r IP MAC; do
                for sel in $SELECTIONS; do
                    if [ "$sel" -eq "$i" ]; then
                        CID=$(echo $IP | awk -F. '{printf "%d%02d\n",$3,$4}')
                        CID_DL="1:${CID}0"
                        CID_UL="1:${CID}00"

                        tc class add dev $BR_IFACE parent 1:1 classid $CID_DL htb rate ${DL_KBPS}kbit ceil ${DL_KBPS}kbit 2>/dev/null
                        tc class add dev $BR_IFACE parent 1:1 classid $CID_UL htb rate ${UL_KBPS}kbit ceil ${UL_KBPS}kbit 2>/dev/null

                        tc filter add dev $BR_IFACE protocol ip parent 1:0 prio 1 u32 match ip dst $IP flowid $CID_DL 2>/dev/null
                        tc filter add dev $BR_IFACE protocol ip parent 1:0 prio 1 u32 match ip src $IP flowid $CID_UL 2>/dev/null

                        echo "[完成] 已限制 $IP 上行 ${UL_MB}MB/s, 下行 ${DL_MB}MB/s"
                    fi
                done
                i=$((i+1))
            done
            ;;
        2)
            delete_limits
            ;;
        3)
            show_limits
            ;;
        4)
            echo "退出管理"
            break
            ;;
        *)
            echo "[提示] 无效选择，请重新输入"
            ;;
    esac
done
